import os
import akshare as ak
import pandas as pd
import numpy as np
from utils import FileUtils

class MarketAnalysisFlow:
    def __init__(self, output_data_dir=None, output_html_path=None):
        """初始化市场分析流程.

        Args:
            output_data_dir (str, optional): 保存Excel数据文件的目录路径.
                                            默认为脚本所在目录下的 "data" 文件夹.
            output_html_path (str, optional): 生成HTML热力图的完整文件路径.
                                             默认为脚本所在目录下的 "board_heat_trend.html".
        """
        self.calculator = BoardHeatScoreCalculator()
        
        script_dir = os.path.dirname(os.path.abspath(__file__))
        
        self.data_dir = output_data_dir if output_data_dir is not None else os.path.join(script_dir, "data")
        self.html_path = output_html_path if output_html_path is not None else os.path.join(script_dir, "board_heat_trend.html")

        # 确保输出目录存在
        if not os.path.exists(self.data_dir):
            os.makedirs(self.data_dir, exist_ok=True)
        
        html_output_parent_dir = os.path.dirname(self.html_path)
        if not os.path.exists(html_output_parent_dir):
            os.makedirs(html_output_parent_dir, exist_ok=True)

    def calculate_board_heat_scores(self):
        """计算板块热力得分并保存"""
        heat_scores = self.calculator.run()
        file_path = FileUtils.save_to_excel(heat_scores, self.data_dir)
        print(f"板块热力得分已保存到: {file_path}")
        return file_path

    def generate_board_heat_trend_heatmap(self):
        """生成板块热力趋势热图"""
        # 设置颜色方案：从深蓝色(0)到正红色(100)
        colorscale = [
            [0, 'rgb(0, 0, 128)'],      # 深蓝色
            [0.25, 'rgb(0, 100, 255)'],  # 蓝色
            [0.5, 'rgb(255, 255, 255)'], # 白色
            [0.75, 'rgb(255, 100, 0)'],  # 橙色
            [1, 'rgb(255, 0, 0)']        # 正红色
        ]
        
        html_file = FileUtils.generate_heatmap_html(self.data_dir, self.html_path, colorscale=colorscale)
        
        if html_file:
            print(f"热力图已生成: {html_file}")
            return html_file
        else:
            print("热力图生成失败")
            return None

    def run_analysis(self):
        """执行完整的分析流程"""
        print("开始计算板块热力得分...")
        self.calculate_board_heat_scores()
        print("\n开始生成板块热力趋势热图...")
        self.generate_board_heat_trend_heatmap()
        print("\n市场分析流程执行完毕。")

class BoardHeatScoreCalculator:
    def __init__(self):
        """初始化板块热力得分计算器"""
        # 权重配置
        self.weights = {
            '涨跌幅': 0.25,
            '成交量': 0.15,
            '成交额': 0.15,
            '净流入': 0.20,
            '上涨比例': 0.15,
            '领涨股': 0.10
        }
    
    def fetch_data(self):
        """获取板块数据"""
        return ak.stock_board_industry_summary_ths()
    
    def preprocess_data(self, df):
        """数据预处理"""
        # 处理涨跌幅
        if df['涨跌幅'].dtype == 'object':
            df['涨跌幅'] = df['涨跌幅'].str.replace('%', '').astype(float)
        elif pd.api.types.is_numeric_dtype(df['涨跌幅']):
            pass
        
        # 处理领涨股-涨跌幅
        if df['领涨股-涨跌幅'].dtype == 'object':
            df['领涨股-涨跌幅'] = df['领涨股-涨跌幅'].str.replace('%', '').astype(float)
        elif pd.api.types.is_numeric_dtype(df['领涨股-涨跌幅']):
            pass
        
        return df
    
    def calculate_normalized_scores(self, df):
        """计算各指标的归一化得分"""
        # 1. 涨跌幅归一化
        df['涨跌幅得分'] = (df['涨跌幅'] + 10) / 20
        df['涨跌幅得分'] = df['涨跌幅得分'].clip(0, 1)
        
        # 2. 成交量归一化
        df['成交量得分'] = df['总成交量'].rank(pct=True)
        
        # 3. 成交额归一化
        df['成交额得分'] = df['总成交额'].rank(pct=True)
        
        # 4. 净流入归一化
        max_inflow = df['净流入'].max()
        min_inflow = df['净流入'].min()
        range_inflow = max_inflow - min_inflow
        df['净流入得分'] = (df['净流入'] - min_inflow) / range_inflow if range_inflow != 0 else 0.5
        
        # 5. 上涨家数比例
        df['上涨比例'] = df['上涨家数'] / (df['上涨家数'] + df['下跌家数'])
        
        # 6. 领涨股表现
        df['领涨股得分'] = (df['领涨股-涨跌幅'] + 10) / 20
        df['领涨股得分'] = df['领涨股得分'].clip(0, 1)
        
        return df
    
    def calculate_heat_score(self, df):
        """计算热力得分"""
        df['热力得分'] = (
            df['涨跌幅得分'] * self.weights['涨跌幅'] +
            df['成交量得分'] * self.weights['成交量'] +
            df['成交额得分'] * self.weights['成交额'] +
            df['净流入得分'] * self.weights['净流入'] +
            df['上涨比例'] * self.weights['上涨比例'] +
            df['领涨股得分'] * self.weights['领涨股']
        ) * 100
        
        # 确保得分在0-100之间
        df['热力得分'] = df['热力得分'].clip(0, 100).round(2)
        
        return df
    
    def get_result(self, df):
        """获取最终结果"""
        return df[['板块', '涨跌幅', '总成交量', '总成交额', '净流入', 
                  '上涨家数', '下跌家数', '热力得分']].sort_values('热力得分', ascending=False)
    
    def run(self):
        """执行完整的计算流程"""
        df = self.fetch_data()
        df = self.preprocess_data(df)
        df = self.calculate_normalized_scores(df)
        df = self.calculate_heat_score(df)
        return self.get_result(df)

if __name__ == "__main__":
    # 示例：使用默认路径
    # analysis_flow = MarketAnalysisFlow()

    # 示例：使用自定义路径
    # script_dir = os.path.dirname(os.path.abspath(__file__))
    # custom_data_dir = os.path.join(script_dir, "custom_data_output")
    # custom_html_path = os.path.join(script_dir, "custom_reports", "custom_heatmap.html")
    # analysis_flow = MarketAnalysisFlow(output_data_dir=custom_data_dir, output_html_path=custom_html_path)

    # 默认情况下，仍然使用无参数构造，将使用类定义中的默认路径
    analysis_flow = MarketAnalysisFlow()
    analysis_flow.run_analysis() 